use std::collections::HashMap;

use container_core::workflow::{SimpleWorkflowEngine, SomeChainsBuilder};
use model_graph_types::{
    assistant::prompts::{
        PromptGenerateOptions, PromptGenerator, PromptGeneratorBuildInTemplate, PromptResult,
        PromptResultBuilder, CODE_CLASS_GENERATE_JSON_TEMPLATE,
        WORKFLOW_RULE_CONFIG_PROMPT_GENERATE_TEMPLATE,
    },
    container::workflow::{WorkflowEngine, WorkflowEngineRunOptionsBuilder},
};

#[derive(Debug, Default)]
pub struct SimplePromptGenerator {}

#[async_trait::async_trait]
impl PromptGenerator for SimplePromptGenerator {
    /// 根据数据库获取到元数据
    async fn generate_with_buildin(
        &self,
        buildin_template: PromptGeneratorBuildInTemplate,
        instruction: &String,
        options: PromptGenerateOptions,
    ) -> anyhow::Result<PromptResult> {
        let template = match buildin_template {
            PromptGeneratorBuildInTemplate::Rule => {
                //
                let template = WORKFLOW_RULE_CONFIG_PROMPT_GENERATE_TEMPLATE;
                Ok(template)
            }
            PromptGeneratorBuildInTemplate::CodeClass => {
                //
                let template = CODE_CLASS_GENERATE_JSON_TEMPLATE;
                Ok(template)
            }
            _ => {
                tracing::error!("buildin_template无效!");
                Err(anyhow::anyhow!("buildin_template无效!"))
            }
        }?;
        //调用大模型，生成提示词
        let chain = SomeChainsBuilder::default()
            .build()?
            .get_rule_block_chain(&String::from(template), instruction)
            .await?;
        // tracing::debug!("{:?}", chain);
        let engine = SimpleWorkflowEngine::default();
        //执行
        let m_id = engine
            .run(
                chain,
                HashMap::new(),
                WorkflowEngineRunOptionsBuilder::default().build()?,
            )
            .await?;
        //获取结果
        let result = engine.result_with_timeout(&m_id, true, 60000_u64).await;
        match result {
            Ok(v) => {
                tracing::debug!("{:?}", v);
                if let Some(text) = v.outputs.get(&String::from("text")) {
                    match text {
                        serde_json::Value::String(t) => {
                            Ok(PromptResultBuilder::default().prompt(t).build()?)
                        }
                        _ => Err(anyhow::anyhow!("返回的结果不是字符串")),
                    }
                } else {
                    Err(anyhow::anyhow!("没有找到结果"))
                }
            }
            Err(err) => Err(anyhow::anyhow!("{}", err)),
        }
    }
}
